/*
 * EquilibriumSolver.cpp
 *
 *  Created on: 13 Jun 2011
 *      Author: allan
 */

#include "EquilibriumSolver.h"

// C++ includes
#include <cmath>
#include <functional>
using namespace std;
using namespace std::placeholders;

// Boost includes
#include <boost/foreach.hpp>
using namespace boost;

// Eigen includes
#include <Eigen/Core>
using namespace Eigen;

// GeoReact includes
#include "MultirootSolver.h"
#include "Utils.h"

EquilibriumSolver::EquilibriumSolver() :
tolerance(1.0E-6), maxiterations(20), numConditions(0)
{}

EquilibriumSolver::EquilibriumSolver(const System& system, const StringVector& equilibriumSpecies) :
system(system), tolerance(1.0E-6), maxiterations(20), numConditions(0)
{
	// Set the indexes of the primary and equilibrium species
	jSpecies = system[system.GetPrimarySpecies()];
	eSpecies = system[equilibriumSpecies];
	
	// Store the equilibrium-controlled reactions which has as the reactants the names of the given equilibrium species
	BOOST_FOREACH(const Index& index, eSpecies)
		reactions.push_back( system.GetReactions()[index] );
	
	// Assemble the reactive (stoichiometric) matrix of the set of equilibrium-controlled reactions
	vej = AssembleReactiveMatrix(reactions, system.GetPrimarySpecies());
	
	// Set the number of primary and equilibrium species
	Nj = jSpecies.size();
	Ne = eSpecies.size();

	alpha.resize(Nj, Nj);
	 beta.resize(Nj, Ne);
	theta.resize(Nj, Nj);
	   mu.resize(Nj, Ne);

	hfunctions.resize(Nj);
}

const bool EquilibriumSolver::SpecifySpeciesTotalMoles(const String& speciesName, double val)
{
	// Check if it is possible to specify another speciation condition
	if(numConditions >= Nj) return false;
	
	// The global index of the given species
	const Index speciesIndex = system[speciesName];
	
	// Try to find the global index of the given species (speciesIndex) in the set of indexes of primary species (jSpecies) and equilibrium species (eSpecies)
	const Index j = std::find(jSpecies.begin(), jSpecies.end(), speciesIndex) - jSpecies.begin();
	const Index e = std::find(eSpecies.begin(), eSpecies.end(), speciesIndex) - eSpecies.begin();
	
	// Check if j is within limits [0, Nj), where Nj is the number of primary species
	if(j < Nj)
	{
		// Set the alpha entry whose row is the current number of speciation conditions (numConditions) and whose column is the local index of the primary species (j)
		alpha(numConditions, j) = 1.0;
		
		// Set the beta entries whose row is the current number of speciation conditions (numConditions) and whose columns are the local indexes of the equilibrium species (e)
		for(unsigned e = 0; e < Ne; ++e)
			beta(numConditions, e) = vej(e, j);
	}
	
	// Check if e is within limits [0, Ne), where Ne is the number of equilibrium species
	if(e < Ne)
		beta(numConditions, e) = 1.0; // Set the respective beta entry
	
	// Set the h function of the current speciation condition as a constant function
	hfunctions[numConditions] = Constant(val);
	
	// Check if at least one of the indexes, j or e, was within limits
	const bool successful = (j < Nj || e < Ne);
	
	// Update the number of specified conditions if this specification succeeded
	if(successful) ++numConditions;
	
	// Return true if the given species was a primary or equilibrium species
	return successful;
}

const bool EquilibriumSolver::SpecifySpeciesMoles(const String& speciesName, double val)
{
	// Check if it is possible to specify another speciation condition
	if(numConditions >= Nj) return false;
	
	// The global index of the given species
	const Index speciesIndex = system[speciesName];
	
	// Try to find the global index of the given species (speciesIndex) in the set of indexes of primary species (jSpecies) and equilibrium species (eSpecies)
	const Index j = std::find(jSpecies.begin(), jSpecies.end(), speciesIndex) - jSpecies.begin();
	const Index e = std::find(eSpecies.begin(), eSpecies.end(), speciesIndex) - eSpecies.begin();
	
	// Check which one, j or e, is within limits, i.e. j in [0, Nj) or e in [0, Ne), where Nj and Ne are the number of primary and equilibrium species respectively
	if(j < Nj) alpha(numConditions, j) = 1.0; // Set the respective alpha entry
	if(e < Ne)  beta(numConditions, e) = 1.0; // Set the respective beta entry
	
	// Set the h function of the current speciation condition as a constant function
	hfunctions[numConditions] = Constant(val);
	
	// Check if at least one of the indexes, j or e, was within limits
	const bool successful = (j < Nj || e < Ne);
	
	// Update the number of specified conditions if this specification succeeded
	if(successful) ++numConditions;
	
	// Return true if the given species was a primary or equilibrium species
	return successful;
}

const bool EquilibriumSolver::SpecifySpeciesActivity(const String& speciesName, double val)
{
	// Check if it is possible to specify another speciation condition
	if(numConditions >= Nj) return false;
	
	// The global index of the given species
	const Index speciesIndex = system[speciesName];
	
	// Try to find the global index of the given species (speciesIndex) in the set of indexes of primary species (jSpecies) and equilibrium species (eSpecies)
	const Index j = std::find(jSpecies.begin(), jSpecies.end(), speciesIndex) - jSpecies.begin();
	const Index e = std::find(eSpecies.begin(), eSpecies.end(), speciesIndex) - eSpecies.begin();
	
	// Check which one, j or e, is within limits, i.e. j in [0, Nj) or e in [0, Ne), where Nj and Ne are the number of primary and equilibrium species respectively
	if(j < Nj) theta(numConditions, j) = 1.0; // Set the respective theta entry
	if(e < Ne)    mu(numConditions, e) = 1.0; // Set the respective mu entry
	
	// Set the h function of the current speciation condition as a constant function
	hfunctions[numConditions] = Constant(val);
	
	// Check if at least one of the indexes, j or e, was within limits
	const bool successful = (j < Nj || e < Ne);
	
	// Update the number of specified conditions if this specification succeeded
	if(successful) ++numConditions;
	
	// Return true if the given species was a primary or equilibrium species
	return successful;
}

const bool EquilibriumSolver::SpecifyAcidity(double pH)
{
	return SpecifySpeciesActivity("H[+]", std::exp(-pH));
}

const double AuxGasPartialPressure(double T, double P, const DoubleVector& n, const DoubleVector& a, const IndexVector& gasesIndexes, double Pg)
{
	// The total number of moles of gaseous species
	double nT = 0.0; for(Index i = 0; i < gasesIndexes.size(); ++i) nT += n[i];

	return nT * Pg/P;
}

const bool EquilibriumSolver::SpecifyGasPartialPressure(const String& gasName, double Pg)
{
	// Find the gaseous phase among all the phases in the reactive system
	const PhasePtr& gaseousPhase = 
			*std::find_if(system.GetPhases().begin(), system.GetPhases().end(), 
					[](const PhasePtr& phase) { return phase->GetName() == "Gaseous"; } );
	
	// The indexes of all the gaseous species
	IndexVector gasesIndexes;
	
	// Loop over all the gaseous species and build gasesIndexes
	BOOST_FOREACH(const String& gasName, gaseousPhase->GetSpecies())
		gasesIndexes.push_back(system[gasName]);
	
	// Reuse the method SpecifySpeciesMoles to set up the alpha entries of this speciation condition (the alpha matrix has the same structure)
	const bool successful = SpecifySpeciesMoles(gasName, 0.0);
	
	// Create the function h for the partial pressure of the gas and overwrite the constant function set up by the call SpecifySpeciesMoles (note that numConditions has been increased in one if the call succeeded)
	if(successful) hfunctions[numConditions - 1] = std::bind(AuxGasPartialPressure, _1, _2, _3, _4, gasesIndexes, Pg);
	
	return successful;
}

const bool EquilibriumSolver::SpecifyChargeBalance()
{
	// Check if it is possible to specify another speciation condition
	if(numConditions >= Nj) return false;
	
	const StringVector& species = system.GetSpecies();
	
	for(Index j = 0; j < Nj; ++j)
		alpha(numConditions, j) = ElectricalCharge(species[jSpecies[j]]);
	
	for(Index e = 0; e < Ne; ++e)
		beta(numConditions, e) = ElectricalCharge(species[eSpecies[e]]);
	
	hfunctions[numConditions] = Constant(0.0);
	
	// Update the number of specified conditions
	++numConditions;
	
	return true;
}

void EquilibriumSolver::ResetConditions()
{
	numConditions = 0;
}

void EquilibriumSolver::SetIndependentVector(const DoubleVector& h)
{
	for(uint j = 0; j < Nj; ++j)
		hfunctions[j] = Constant(h[j]);
}

void EquilibriumSolver::SetSolverOptions(double tolerance, uint maxiterations)
{
	this->tolerance = tolerance;
	this->maxiterations = maxiterations;
}

const DoubleVector EquilibriumSolver::Solve(double T, double P, const DoubleVector& n0)
{
	// The number of unknowns
	const uint numUnknowns = Nj + Ne;
	
	// The vector njne = [nj ne]
	VectorXd njne(numUnknowns);
	
	// Initialise the vector njne with the values from the initial guess n0
	for(uint j = 0; j < Nj; ++j) njne[j]      = n0[jSpecies[j]];
	for(uint e = 0; e < Ne; ++e) njne[e + Nj] = n0[eSpecies[e]];
	
	// Create an equilibrium problem object
	EquilibriumProblem problem(*this, T, P, n0);
	
	// Create a multiroot solver object
	MultirootSolver solver;
	
	// Solve the equilibrium problem
	solver.Solve(problem, njne);
	
	// The solution vector (n) containing the with the (njne) part updated
	DoubleVector n(n0);
	
	// Transfer the solution (njne) to the vector (n)
	for(uint j = 0; j < Nj; ++j) n[jSpecies[j]] = njne[j];
	for(uint e = 0; e < Ne; ++e) n[eSpecies[e]] = njne[e + Nj];
	
	// Return the mole vector (n)
	return n;
}

EquilibriumProblem::EquilibriumProblem(const EquilibriumSolver& es, double T, double P, const DoubleVector& n0) :
es(es), T(T), P(P), n0(n0), Nj(es.Nj), Ne(es.Ne), n(n0), nj(Nj), ne(Ne), a(n0), aj(Nj), ae(Ne), Ke(Ne), Qe(Ne), h(Nj),
Fj(Nj), Fe(Ne), Jjj(Nj, Nj), Jje(Nj, Ne), Jej(Ne, Nj), Jee(Ne, Ne)
{}

void EquilibriumProblem::Function(const VectorXd& njne, VectorXd& F)
{
	// Update the auxiliary data members of the equilibrium problem
	UpdateAuxiliarData(njne);
	
	// Transfer the data from Fj the vector F
	for(uint j = 0; j < Nj; ++j) F[j] = Fj[j];
	
	// Transfer the data from Fe to the vector F
	for(uint e = 0; e < Ne; ++e) F[e + Nj] = Fe[e];
}

void EquilibriumProblem::Jacobian(const VectorXd& njne, MatrixXd& J)
{
	// Update the auxiliary data members of the equilibrium problem
	UpdateAuxiliarData(njne);
	
	// Transfer the data from Jjj to the matrix J
	for(uint row = 0; row < Nj; ++row)
		for(uint col = 0; col < Nj; ++col)
			J(row, col) = Jjj(row, col);
	
	// Transfer the data from Jje to the matrix J
	for(uint row = 0; row < Nj; ++row)
		for(uint col = 0; col < Ne; ++col)
			J(row, col + Nj) = Jje(row, col);
	
	// Transfer the data from Jej to the matrix J
	for(uint row = 0; row < Ne; ++row)
		for(uint col = 0; col < Nj; ++col)
			J(row + Nj, col) = Jej(row, col);
	
	// Transfer the data from Jee to the matrix J
	for(uint row = 0; row < Ne; ++row)
		for(uint col = 0; col < Ne; ++col)
			J(row + Nj, col + Nj) = Jee(row, col);
}

void EquilibriumProblem::PostProcess(VectorXd& njne)
{
	UnderRelax(njne, njneOld);
	
	njneOld = njne;
}

void EquilibriumProblem::UpdateAuxiliarData(const VectorXd& njne)
{
	// Update the (nj) and (ne) part of the vector n
	for(uint j = 0; j < Nj; ++j) n[es.jSpecies[j]] = njne[j];
	for(uint e = 0; e < Ne; ++e) n[es.eSpecies[e]] = njne[e + Nj];
	
	// Update the activity vector a with temperature (t), pressure (P) and the mole vector (n)
	a = es.system.Activities(T, P, n);
	
	// Update the mole vectors nj and ne
	for(uint j = 0; j < Nj; ++j) nj[j] = n[es.jSpecies[j]];
	for(uint e = 0; e < Ne; ++e) ne[e] = n[es.eSpecies[e]];
	
	// Update the activity vectors aj and ae
	for(uint j = 0; j < Nj; ++j) aj[j] = a[es.jSpecies[j]];
	for(uint e = 0; e < Ne; ++e) ae[e] = a[es.eSpecies[e]];
	
	// Update the vector h (the right-hand side values of the speciation conditions)
	for(uint j = 0; j < Nj; ++j) h[j]  = es.hfunctions[j](T, P, n, a);
	
	// Update the equilibrium constants of the equilibrium reactions
	for(uint e = 0; e < Ne; ++e) Ke[e] = es.reactions[e]->EquilibriumConstant(T, P);
	
	// Update the reaction quotients of the equilibrium reactions
	for(uint e = 0; e < Ne; ++e) Qe[e] = es.reactions[e]->ReactionQuotient(a);
	
	// Update the Fj vector
	Fj = (es.alpha * nj) + (es.beta * ne) + (es.theta * aj) + (es.mu * ae) - h;
	
	// Update the Fe vector
	Fe = (Ke.array()/Qe.array()).log(); // component-wise division and application of log
	
	// Update the Jjj matrix
	for(uint row = 0; row < Nj; ++row)
		for(uint col = 0; col < Nj; ++col)
			Jjj(row, col) = es.alpha(row, col) + es.theta(row, col) * aj[col]/nj[col];
	
	// Update the Jje matrix
	for(uint row = 0; row < Nj; ++row)
		for(uint col = 0; col < Ne; ++col)
			Jje(row, col) = es.beta(row, col) + es.mu(row, col) * ae[col]/ne[col];
	
	// Update the Jej matrix
	for(uint row = 0; row < Ne; ++row)
		for(uint col = 0; col < Nj; ++col)
			Jej(row, col) = - es.vej(row, col)/nj[col];
	
	// Update the Jee matrix
	for(uint row = 0; row < Ne; ++row)
			Jee(row, row) = 1.0/ne[row];
}

void EquilibriumProblem::UnderRelax(VectorXd& njne, const VectorXd& njneOld)
{
	for(uint i = 0; i < uint(njne.rows()); ++i)
	{
		const double delta  = njne[i] - njneOld[i];
		const double factor = max(1.0, -2.0 * delta/njneOld[i]);
		
		njne[i] = njneOld[i] + delta/factor;
	}
}
